{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.Tensor'>\n",
      "torch.int64\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "# 假设probabilities是一个包含概率分布的tensor\n",
    "probabilities = torch.tensor([[0.1, 0.8, 0.1], [0.3, 0.4, 0.3]], dtype=torch.float)\n",
    "# 找到最大概率的索引\n",
    "indices = probabilities.argmax(dim=1)\n",
    "print(type(indices))\n",
    "print(indices.dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1380, 512])\n",
      "tensor([[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         ...,\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "         [0., 0., 0.,  ..., 0., 0., 0.]]])\n",
      "action_logp: torch.Size([1, 11, 512])\n",
      "action_token: torch.Size([1, 11])\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_19136/3672849072.py:30: UserWarning: Implicit dimension choice for softmax has been deprecated. Change the call to include dim=X as an argument.\n",
      "  action_logp =  F.softmax(action_logits)\n"
     ]
    }
   ],
   "source": [
    "# torch.is_same_size\n",
    "\n",
    "import torch\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "seqlen=10\n",
    "\n",
    "pos = np.arange(0, seqlen, 1)\n",
    "pos=torch.tensor(pos)\n",
    "pos = F.one_hot(pos, num_classes=seqlen)\n",
    "pos.unsqueeze(0).shape\n",
    "\n",
    "from torch import nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "output_logits =torch.zeros([1, 1380, 512])\n",
    "print(output_logits.shape)\n",
    "\n",
    "seqlen=15\n",
    "time_step_tokens=92\n",
    "num_image_tokens=81\n",
    "output_logits = np.reshape(\n",
    "        output_logits, (1, seqlen, time_step_tokens, -1)\n",
    "    )\n",
    "action_logits = output_logits[:, -1, ...]\n",
    "action_logits = action_logits[:, num_image_tokens - 1 : -1]\n",
    "print(action_logits)\n",
    "action_logp =  F.softmax(action_logits)\n",
    "print('action_logp:',action_logp.shape)\n",
    "action_token = torch.argmax(action_logp, axis=-1)\n",
    "print('action_token:',action_token.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1380, 512])\n",
      "torch.Size([1, 1380, 512])\n"
     ]
    }
   ],
   "source": [
    "x1 =torch.zeros([1, 1380, 512])\n",
    "attn_mask =torch.zeros([1380, 1380])\n",
    "print(x1.shape)\n",
    "num_heads=8\n",
    "dropout_rate=0.1\n",
    "\n",
    "if 1: \n",
    "    multihead_attn = nn.MultiheadAttention(\n",
    "        batch_first=True,\n",
    "        embed_dim = x1.shape[-1],\n",
    "        num_heads=num_heads,\n",
    "        dropout=dropout_rate,\n",
    "    )\n",
    "    x1, attn_output_weights = multihead_attn(x1, x1, x1,attn_mask=attn_mask)\n",
    "    print(x1.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "# 设定输入序列长度和头数\n",
    "seq_len = 5\n",
    "num_heads = 2\n",
    "\n",
    "# 创建模拟输入\n",
    "input = torch.randn(1, seq_len, 256)\n",
    "\n",
    "# 创建模拟的att_mask，将第2和第3个位置屏蔽掉\n",
    "att_mask = torch.zeros(seq_len, seq_len)\n",
    "att_mask[1:3, :] = -1e9  # 设置为一个较小的值\n",
    "att_mask[:, 1:3] = -1e9  # 设置为一个较小的值\n",
    "\n",
    "# 创建多头注意力机制实例\n",
    "attention = torch.nn.MultiheadAttention(256, num_heads,batch_first=True)\n",
    "\n",
    "# 使用att_mask计算注意力权重\n",
    "output, attn_weights = attention(input, input, input, attn_mask=att_mask)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "model_scope",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
